#include <iostream>
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include "chrono"
#include "vector"
#include "fstream"
#include "utils.h"
#include "math.h"
#include "opencv2/opencv.hpp"
#include "sstream"
#include "string"

static Logger gLogger;


static const int INPUT_H = 640;
static const int INPUT_W = 640;
static const int OUTPUT_SIZE = INPUT_W*INPUT_H;

using namespace cv;
using namespace std;
using namespace nvinfer1;
const char *INPUT_BLOB_NAME = "data";
const char *OUTPUT_BLOB_NAME = "D";

vector<float> loadWeights(const string &filePath) {
    int size = 0;
    ifstream file(filePath, ios::binary);
    file.read((char *) &size, 4);
    char *floatWeight = new char[size * 4];
    float *fp = (float *) floatWeight;
    file.read(floatWeight, 4 * size);
    vector<float> weights(fp, fp + size);
    delete[] floatWeight;
    file.close();
    return weights;
}

IConvolutionLayer *AddConv(INetworkDefinition *network,string convpath, ITensor &input, int out_channel, int hstride = 1, int wstride = 1,
                          int pad = 1, int kernel_size = 3,int dila =1,int group=1){
    vector<float> conv_weights;
    conv_weights= loadWeights(convpath);
    int w_size=conv_weights.size();
    Weights convWeights{nvinfer1::DataType::kFLOAT, nullptr, w_size};
    Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, 0};
    float *val_wt = new float[w_size];
    for (int i = 0; i < w_size; i++) {
        val_wt[i] = conv_weights[i];
    }
    convWeights.values = val_wt;
    IConvolutionLayer *conv1 = network->addConvolutionNd(input, out_channel, DimsHW{kernel_size, kernel_size},
                                                         convWeights, convBias);
    assert(conv1);
    conv1->setStrideNd(DimsHW{hstride, wstride});
    conv1->setPaddingNd(DimsHW{pad, pad});
    conv1->setDilationNd(DimsHW{dila,dila});
    conv1->setNbGroups(group);
    return conv1;
}


IDeconvolutionLayer *AddDenConv(INetworkDefinition *network,string weightpath, string biaspath,ITensor &input, int out_channel, int hstride = 1, int wstride = 1,
                           int pad = 1, int kernel_size = 3,int dila =1,int group=1){
    vector<float> conv_weights;
    conv_weights= loadWeights(weightpath);
    int w_size=conv_weights.size();

    vector<float> conv_bias;
    conv_bias= loadWeights(biaspath);


    Weights convWeights{nvinfer1::DataType::kFLOAT, nullptr, w_size};
    Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, out_channel};
    float *val_wt = new float[w_size];
    for (int i = 0; i < w_size; i++) {
        val_wt[i] = conv_weights[i];
    }
    convWeights.values = val_wt;


    float *val_bias=new float [out_channel];
    for(int i=0;i<out_channel;i++){
        val_bias[i] = 0.0;
        if(conv_bias.size() != 0){
            val_bias[i] = conv_bias[i];
        }
    }
    convBias.values = val_bias;
    IDeconvolutionLayer *dencov=network->addDeconvolutionNd(input,out_channel,DimsHW{kernel_size,kernel_size},convWeights,convBias);

    assert(dencov);
    dencov->setStrideNd(DimsHW{hstride, wstride});
    dencov->setPaddingNd(DimsHW{pad, pad});
    dencov->setNbGroups(group);
    return dencov;
}

IScaleLayer *AddBN(INetworkDefinition *network,string bnpath, ITensor &input){
    float eps = 1e-5;
    vector<float> bn_1_bias;
    vector<float> bn_1_mean;
    vector<float> bn_1_var;
    vector<float> bn_1_weight;
//    cout<<bnpath<<endl;

    bn_1_bias = loadWeights(bnpath+"bn.bias.wgt");
    bn_1_mean = loadWeights(bnpath+"bn.running_mean.wgt");
    bn_1_var = loadWeights(bnpath+"bn.running_var.wgt");
    bn_1_weight = loadWeights(bnpath+"bn.weight.wgt");
    int bn_size = bn_1_weight.size();

    float *scval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        scval[i] = bn_1_weight[i] / sqrt(bn_1_var[i] + eps);
    }

    float *shval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        shval[i] = bn_1_bias[i] - bn_1_mean[i] * bn_1_weight[i] / sqrt(bn_1_var[i] + eps);
    }
    float *pval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        pval[i] = 1.0;
    }
    Weights shift{nvinfer1::DataType::kFLOAT, shval, bn_size};
    Weights scale{nvinfer1::DataType::kFLOAT, scval, bn_size};
    Weights power{nvinfer1::DataType::kFLOAT, pval, bn_size};

    IScaleLayer *scale1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
    assert(scale1);
    return scale1;
}
IScaleLayer *AddBN2(INetworkDefinition *network,string bnpath, ITensor &input){
    float eps = 1e-5;
    vector<float> bn_1_bias;
    vector<float> bn_1_mean;
    vector<float> bn_1_var;
    vector<float> bn_1_weight;
//    cout<<bnpath<<endl;

    bn_1_bias = loadWeights(bnpath+"bias.wgt");
    bn_1_mean = loadWeights(bnpath+"running_mean.wgt");
    bn_1_var = loadWeights(bnpath+"running_var.wgt");
    bn_1_weight = loadWeights(bnpath+"weight.wgt");
    int bn_size = bn_1_weight.size();
    float *scval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        scval[i] = bn_1_weight[i] / sqrt(bn_1_var[i] + eps);
    }

    float *shval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        shval[i] = bn_1_bias[i] - bn_1_mean[i] * bn_1_weight[i] / sqrt(bn_1_var[i] + eps);
    }
    float *pval = reinterpret_cast<float *>(malloc(sizeof(float) * bn_size));
    for (int i = 0; i < bn_size; i++) {
        pval[i] = 1.0;
    }
    Weights shift{nvinfer1::DataType::kFLOAT, shval, bn_size};
    Weights scale{nvinfer1::DataType::kFLOAT, scval, bn_size};
    Weights power{nvinfer1::DataType::kFLOAT, pval, bn_size};

    IScaleLayer *scale1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power);
    assert(scale1);
    return scale1;
}

IPluginV2Layer *AddHSwish(INetworkDefinition *network,ITensor *input){
    auto creator = getPluginRegistry()->getPluginCreator("HardSwishLayer_TRT", "1");
    const PluginFieldCollection* pluginData = creator->getFieldNames();
    IPluginV2 *pluginObj = creator->createPlugin("hardswish", pluginData);
    ITensor* inputTensors[] = {input};
    return network->addPluginV2(inputTensors, 1, *pluginObj);

}
IResizeLayer *AddUPSample2X(INetworkDefinition *network ,ITensor &input,int factor=2){
    auto *upsample1=network->addResize(input);
    Dims3 dims;
    dims.d[0] = input.getDimensions().d[0];
    dims.d[1] = input.getDimensions().d[1]*factor;
    dims.d[2] = input.getDimensions().d[2]*factor;

    upsample1->setOutputDimensions(dims);
    upsample1->setResizeMode(ResizeMode::kNEAREST);
    return upsample1;
}


ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, nvinfer1::DataType dt) {
    INetworkDefinition *network = builder->createNetworkV2(0U);
    //x
#pragma region input
    ITensor *data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W});
    assert(data);
#pragma endregion

#pragma region backbone
    //conv1
    auto *conv1= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.conv1.conv.weight.wgt",*data,8,2,2,1,3);
    assert(conv1);
    auto *bn1= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.conv1.",*conv1->getOutput(0));
    bn1->setName("bn1");
    assert(bn1);
    auto *hswish= AddHSwish(network,bn1->getOutput(0));
    //stage0
    assert(hswish);
    auto *conv2= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv0.conv.weight.wgt",*hswish->getOutput(0),8,1,1,0,1);
    assert(conv2);
    auto *bn2= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv0.",*conv2->getOutput(0));
    bn2->setName("bn2");
    assert(bn2);
    auto *relu1=network->addActivation(*bn2->getOutput(0),ActivationType::kRELU);
    assert(relu1);
    auto *conv3= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv1.conv.weight.wgt",*relu1->getOutput(0),8,1,1,1,3,1,8);
    assert(conv3);
    auto *bn3= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv1.",*conv3->getOutput(0));
    bn3->setName("bn3");
    assert(bn3);
    auto *relu2=network->addActivation(*bn3->getOutput(0),ActivationType::kRELU);
    assert(relu2);
    auto *conv4= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv2.conv.weight.wgt",*relu2->getOutput(0),8,1,1,0,1,1,1);
    assert(conv4);
    auto bn4 = AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.0.conv2.",*conv4->getOutput(0));
    bn4->setName("bn4");
    assert(bn4);
    auto *elt1=network->addElementWise(*bn4->getOutput(0),*hswish->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt1);
    auto *conv5= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv0.conv.weight.wgt",*elt1->getOutput(0),32,1,1,0,1,1,1);
    assert(conv5);
    auto *bn5= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv0.",*conv5->getOutput(0));
    bn5->setName("bn5");
    assert(bn5);
    auto *relu3=network->addActivation(*bn5->getOutput(0),ActivationType::kRELU);
    assert(relu3);
    auto *conv6= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv1.conv.weight.wgt",*relu3->getOutput(0),32,2,2,1,3,1,32);
    assert(conv6);
    auto *bn6= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv1.",*conv6->getOutput(0));
    assert(bn6);
    auto *relu4=network->addActivation(*bn6->getOutput(0),ActivationType::kRELU);
    assert(relu4);
    auto *conv7= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv2.conv.weight.wgt",*relu4->getOutput(0),16,1,1,0,1,1,1);
    assert(conv7);
    auto *bn7= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.1.conv2.",*conv7->getOutput(0));
    assert(bn7);
    auto *conv8 = AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv0.conv.weight.wgt",*bn7->getOutput(0),40,1,1,0,1,1,1);
    assert(conv8);
    auto *bn8= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv0.",*conv8->getOutput(0));
    assert(bn8);
    auto *relu5=network->addActivation(*bn8->getOutput(0),ActivationType::kRELU);
    assert(relu5);
    auto *conv9= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv1.conv.weight.wgt",*relu5->getOutput(0),40,1,1,1,3,1,40);
    assert(conv9);
    auto *bn9 = AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv1.",*conv9->getOutput(0));
    assert(bn9);
    auto *relu6=network->addActivation(*bn9->getOutput(0),ActivationType::kRELU);
    assert(relu6);
    auto *conv10= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv2.conv.weight.wgt",*relu6->getOutput(0),16,1,1,0,1,1,1);
    assert(conv10);
    auto *bn10= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.0.2.conv2.",*conv10->getOutput(0));
    assert(bn10);
    auto *elt2=network->addElementWise(*bn10->getOutput(0),*bn7->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt2);
    //stage1
    auto *conv11= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv0.conv.weight.wgt",*elt2->getOutput(0),40,1,1,0,1,1,1);
    assert(conv11);
    auto *bn11= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv0.",*conv11->getOutput(0));
    assert(bn11);
    auto *relu7=network->addActivation(*bn11->getOutput(0),ActivationType::kRELU);
    assert(relu7);
    auto *conv12= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv1.conv.weight.wgt",*relu7->getOutput(0),40,2,2,2,5,1,40);
    assert(conv12);
    auto *bn12= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv1.",*conv12->getOutput(0));
    assert(bn12);
    auto *relu8=network->addActivation(*bn12->getOutput(0),ActivationType::kRELU);
    assert(relu8);
    auto *conv13= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv2.conv.weight.wgt",*relu8->getOutput(0),24,1,1,0,1,1,1);
    assert(conv13);
    auto *bn13= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.0.conv2.",*conv13->getOutput(0));
    assert(bn13);
    auto *conv14= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv0.conv.weight.wgt",*bn13->getOutput(0),64,1,1,0,1,1,1);
    assert(conv14);
    auto *bn14= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv0.",*conv14->getOutput(0));
    assert(bn14);
    auto *relu9=network->addActivation(*bn14->getOutput(0),ActivationType::kRELU);
    assert(relu9);
    auto *conv15= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv1.conv.weight.wgt",*relu9->getOutput(0),64,1,1,2,5,1,64);
    assert(conv15);
    auto *bn15= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv1.",*conv15->getOutput(0));
    assert(bn15);
    auto *relu10= network->addActivation(*bn15->getOutput(0),ActivationType::kRELU);
    assert(relu10);
    auto *conv16= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv2.conv.weight.wgt",*relu10->getOutput(0),24,1,1,0,1,1,1);
    assert(conv16);
    auto *bn16= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.1.conv2.",*conv16->getOutput(0));
    assert(bn16);
    auto *elt3=network->addElementWise(*bn16->getOutput(0),*bn13->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt3);
    auto *conv17= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv0.conv.weight.wgt",*elt3->getOutput(0),64,1,1,0,1,1,1);
    assert(conv17);
    auto *bn17= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv0.",*conv17->getOutput(0));
    assert(bn17);
    auto *relu11=network->addActivation(*bn17->getOutput(0),ActivationType::kRELU);
    assert(relu11);
    auto *conv18= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv1.conv.weight.wgt",*relu11->getOutput(0),64,1,1,2,5,1,64);
    assert(conv18);
    auto *bn18= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv1.",*conv18->getOutput(0));
    assert(bn18);
    auto *relu12=network->addActivation(*bn18->getOutput(0),ActivationType::kRELU);
    assert(relu12);
    auto *conv19= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv2.conv.weight.wgt",*relu12->getOutput(0),24,1,1,0,1,1,1);
    assert(conv19);
    auto *bn19= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.1.2.conv2.",*conv19->getOutput(0));
    assert(bn19);
    auto *elt4=network->addElementWise(*bn19->getOutput(0),*elt3->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt4);
    //stage2
    auto *conv20= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv0.conv.weight.wgt",*elt4->getOutput(0),120,1,1,0,1,1,1);
    assert(conv20);
    auto *bn20= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv0.",*conv20->getOutput(0));
    assert(bn20);
    auto *hswish2= AddHSwish(network,bn20->getOutput(0));
    assert(hswish2);
    auto *conv21= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv1.conv.weight.wgt",*hswish2->getOutput(0),120,2,2,1,3,1,120);
    assert(conv21);
    auto *bn21= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv1.",*conv21->getOutput(0));
    assert(bn21);
    auto *hswish3= AddHSwish(network,bn21->getOutput(0));
    assert(hswish3);
    auto *conv22= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv2.conv.weight.wgt",*hswish3->getOutput(0),40,1,1,0,1,1,1);
    assert(conv22);
    auto *bn22= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.0.conv2.",*conv22->getOutput(0));
    assert(bn22);
    auto *conv23= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv0.conv.weight.wgt",*bn22->getOutput(0),104,1,1,0,1,1,1);
    assert(conv23);
    auto *bn23= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv0.",*conv23->getOutput(0));
    assert(bn23);
    auto *hswish4= AddHSwish(network,bn23->getOutput(0));
    assert(hswish4);
    auto *conv24= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv1.conv.weight.wgt",*hswish4->getOutput(0),104,1,1,1,3,1,104);
    assert(conv24);
    auto *bn24= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv1.",*conv24->getOutput(0));
    assert(bn24);
    auto *hswish5= AddHSwish(network,bn24->getOutput(0));
    assert(hswish5);
    auto *conv25= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv2.conv.weight.wgt",*hswish5->getOutput(0),40,1,1,0,1,1,1);
    assert(conv25);
    auto *bn25 = AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.1.conv2.",*conv25->getOutput(0));
    assert(bn25);
    auto *elt5=network->addElementWise(*bn25->getOutput(0),*bn22->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt5);
    auto *conv26= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv0.conv.weight.wgt",*elt5->getOutput(0),96,1,1,0,1,1,1);
    assert(conv26);
    auto *bn26= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv0.",*conv26->getOutput(0));
    assert(bn26);
    auto *hswish6= AddHSwish(network,bn26->getOutput(0));
    assert(hswish6);
    auto *conv27= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv1.conv.weight.wgt",*hswish6->getOutput(0),96,1,1,1,3,1,96);
    assert(conv27);
    auto *bn27= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv1.",*conv27->getOutput(0));
    assert(bn27);
    auto *hswish7= AddHSwish(network,bn27->getOutput(0));
    assert(hswish7);
    auto *conv28= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv2.conv.weight.wgt",*hswish7->getOutput(0),40,1,1,0,1,1,1);
    assert(conv28);
    auto *bn28= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.2.conv2.",*conv28->getOutput(0));
    assert(bn28);
    auto *elt6=network->addElementWise(*bn28->getOutput(0),*elt5->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt6);
    auto *conv29= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv0.conv.weight.wgt",*elt6->getOutput(0),96,1,1,0,1,1,1);
    assert(conv29);
    auto *bn29= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv0.",*conv29->getOutput(0));
    assert(bn29);
    auto *hswish8= AddHSwish(network,bn29->getOutput(0));
    assert(hswish8);
    auto *conv30= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv1.conv.weight.wgt",*hswish8->getOutput(0),96,1,1,1,3,1,96);
    assert(conv30);
    auto *bn30= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv1.",*conv30->getOutput(0));
    assert(bn30);
    auto *hswish9= AddHSwish(network,bn30->getOutput(0));
    assert(hswish9);
    auto *conv31= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv2.conv.weight.wgt",*hswish9->getOutput(0),40,1,1,0,1,1,1);
    assert(conv31);
    auto *bn31= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.3.conv2.",*conv31->getOutput(0));
    assert(bn31);
    auto *elt7=network->addElementWise(*bn31->getOutput(0),*elt6->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt7);
    auto *conv32= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv0.conv.weight.wgt",*elt7->getOutput(0),240,1,1,0,1,1,1);
    assert(conv32);
    auto *bn32= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv0.",*conv32->getOutput(0));
    assert(bn32);
    auto *hswish10= AddHSwish(network,bn32->getOutput(0));
    assert(hswish10);
    auto *conv33= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv1.conv.weight.wgt",*hswish10->getOutput(0),240,1,1,1,3,1,240);
    assert(conv33);
    auto *bn33= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv1.",*conv33->getOutput(0));
    assert(bn33);
    auto *hswish11= AddHSwish(network,bn33->getOutput(0));
    assert(hswish11);
    auto *conv34= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv2.conv.weight.wgt",*hswish11->getOutput(0),56,1,1,0,1,1,1);
    assert(conv34);
    auto *bn34= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.4.conv2.",*conv34->getOutput(0));
    assert(bn34);
    auto *conv35= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv0.conv.weight.wgt",*bn34->getOutput(0),336,1,1,0,1,1,1);
    assert(conv35);
    auto *bn35= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv0.",*conv35->getOutput(0));
    assert(bn35);
    auto *hswish12= AddHSwish(network,bn35->getOutput(0));
    assert(hswish12);
    auto *conv36= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv1.conv.weight.wgt",*hswish12->getOutput(0),336,1,1,1,3,1,336);
    assert(conv36);
    auto *bn36= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv1.",*conv36->getOutput(0));
    assert(bn36);
    auto *hswish13= AddHSwish(network,bn36->getOutput(0));
    assert(hswish13);
    auto *conv37= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv2.conv.weight.wgt",*hswish13->getOutput(0),56,1,1,0,1,1,1);
    assert(conv37);
    auto *bn37= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.2.5.conv2.",*conv37->getOutput(0));
    assert(bn37);
    auto *elt8=network->addElementWise(*bn37->getOutput(0),*bn34->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt8);
    //stage3
    auto *conv38= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv0.conv.weight.wgt",*elt8->getOutput(0),336,1,1,0,1,1,1);
    assert(conv38);
    auto *bn38= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv0.",*conv38->getOutput(0));
    assert(bn38);
    auto *hswish14= AddHSwish(network,bn38->getOutput(0));
    assert(hswish14);
    auto *conv39= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv1.conv.weight.wgt",*hswish14->getOutput(0),336,2,2,2,5,1,336);
    assert(conv39);
    auto *bn39= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv1.",*conv39->getOutput(0));
    assert(bn39);
    auto *hswish15= AddHSwish(network,bn39->getOutput(0));
    assert(hswish15);
    auto *conv40= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv2.conv.weight.wgt",*hswish15->getOutput(0),80,1,1,0,1,1,1);
    assert(conv40);
    auto *bn40= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.0.conv2.",*conv40->getOutput(0));
    assert(bn40);
    auto *conv41= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv0.conv.weight.wgt",*bn40->getOutput(0),480,1,1,0,1,1,1);
    assert(conv41);
    auto *bn41= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv0.",*conv41->getOutput(0));
    assert(bn41);
    auto *hswish16= AddHSwish(network,bn41->getOutput(0));
    assert(hswish16);
    auto *conv42= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv1.conv.weight.wgt",*hswish16->getOutput(0),480,1,1,2,5,1,480);
    assert(conv42);
    auto *bn42= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv1.",*conv42->getOutput(0));
    assert(bn42);
    auto *hswish17= AddHSwish(network,bn42->getOutput(0));
    assert(hswish17);
    auto *conv43= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv2.conv.weight.wgt",*hswish17->getOutput(0),80,1,1,0,1,1,1);
    assert(conv43);
    auto *bn43= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.1.conv2.",*conv43->getOutput(0));
    assert(bn43);
    auto *elt9=network->addElementWise(*bn43->getOutput(0),*bn40->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt9);
    auto *conv44= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv0.conv.weight.wgt",*elt9->getOutput(0),480,1,1,0,1,1,1);
    assert(conv44);
    auto *bn44= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv0.",*conv44->getOutput(0));
    assert(bn44);
    auto *hswish18= AddHSwish(network,bn44->getOutput(0));
    assert(hswish18);
    auto *conv45= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv1.conv.weight.wgt",*hswish18->getOutput(0),480,1,1,2,5,1,480);
    assert(conv45);
    auto *bn45= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv1.",*conv45->getOutput(0));
    assert(bn45);
    auto *hswish19= AddHSwish(network,bn45->getOutput(0));
    assert(hswish19);
    auto *conv46= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv2.conv.weight.wgt",*hswish19->getOutput(0),80,1,1,0,1,1,1);
    assert(conv46);
    auto *bn46= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.2.conv2.",*conv46->getOutput(0));
    assert(bn46);
    auto *elt10=network->addElementWise(*bn46->getOutput(0),*elt9->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt10);
    auto *conv47= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.3.conv.weight.wgt",*elt10->getOutput(0),480,1,1,0,1,1,1);
    assert(conv47);
    auto *bn47= AddBN(network,"/home/luotianhang/CLionProjects/dbnet/wts/backbone.stages.3.3.",*conv47->getOutput(0));
    assert(bn47);
    auto *hswish20= AddHSwish(network,bn47->getOutput(0));
    assert(hswish20);

#pragma endregion



#pragma region neck
    auto *conv48= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.in5_conv.weight.wgt",*hswish20->getOutput(0),96,1,1,0,1,1,1);
    assert(conv48);
    auto *conv49= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.in4_conv.weight.wgt",*elt8->getOutput(0),96,1,1,0,1,1,1);
    assert(conv49);
    auto *conv50= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.in3_conv.weight.wgt",*elt4->getOutput(0),96,1,1,0,1,1,1);
    assert(conv50);
    auto *conv51= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.in2_conv.weight.wgt",*elt2->getOutput(0),96,1,1,0,1,1,1);
    assert(conv51);

    auto *upsampl1= AddUPSample2X(network,*conv48->getOutput(0));

    assert(upsampl1);

    auto *elt11= network->addElementWise(*upsampl1->getOutput(0),*conv49->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt11);
    auto *upsampl2= AddUPSample2X(network,*elt11->getOutput(0));
    assert(upsampl2);
    auto *elt12= network->addElementWise(*upsampl2->getOutput(0),*conv50->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt12);
    auto *upsampl3= AddUPSample2X(network,*elt12->getOutput(0));
    assert(upsampl3);
    auto *elt13= network->addElementWise(*upsampl3->getOutput(0),*conv51->getOutput(0),ElementWiseOperation::kSUM);
    assert(elt13);

    auto *conv52= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.p5_conv.weight.wgt",*conv48->getOutput(0),24,1,1,1,3,1,1);
    assert(conv52);
    auto *conv53= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.p4_conv.weight.wgt",*elt11->getOutput(0),24,1,1,1,3,1,1);
    assert(conv53);
    auto *conv54= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.p3_conv.weight.wgt",*elt12->getOutput(0),24,1,1,1,3,1,1);
    assert(conv54);
    auto *conv55= AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/neck.p2_conv.weight.wgt",*elt13->getOutput(0),24,1,1,1,3,1,1);
    assert(conv55);

    auto upsampel4= AddUPSample2X(network,*conv52->getOutput(0),8);
    assert(upsampel4);
    auto upsampel5= AddUPSample2X(network,*conv53->getOutput(0),4);
    assert(upsampel5);
    auto upsampel6= AddUPSample2X(network,*conv54->getOutput(0),2);
    assert(upsampel6);

    ITensor **b = new ITensor*[4];
    b[0]=upsampel4->getOutput(0);
    b[1]=upsampel5->getOutput(0);
    b[2]=upsampel6->getOutput(0);
    b[3]=conv55->getOutput(0);

    auto *concat=network->addConcatenation(b,4);
    concat->setAxis(0);
    assert(concat);



#pragma endregion

#pragma region head


    auto *conv56 = AddConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv1.weight.wgt",*concat->getOutput(0),24,1,1,1,3,1,1);
    assert(conv56);
    auto *bn56= AddBN2(network,"/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv_bn1.",*conv56->getOutput(0));
    assert(bn56);
    auto *relu13=network->addActivation(*bn56->getOutput(0),ActivationType::kRELU);
    assert(relu13);

    auto *deconv1= AddDenConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv2.weight.wgt",
                              "/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv2.bias.wgt",
                             *relu13->getOutput(0),24,2,2,0,2,1,1 );
    assert(deconv1);

    auto *bn57= AddBN2(network,"/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv_bn2.",*deconv1->getOutput(0));
    bn57->setName("bn57");
    assert(bn57);
    auto *relu14=network->addActivation(*bn57->getOutput(0),ActivationType::kRELU);
    assert(relu14);

    auto *deconv2= AddDenConv(network,"/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv3.weight.wgt","/home/luotianhang/CLionProjects/dbnet/wts/head.binarize.conv3.bias.wgt",
                              *relu14->getOutput(0),1,2,2,0,2,1,1);
    deconv2->setName("deconv2");
    assert(deconv2);
    auto *sigmoid=network->addActivation(*deconv2->getOutput(0),ActivationType::kSIGMOID);
    assert(sigmoid);




#pragma endregion




    sigmoid->getOutput(0)->setName(OUTPUT_BLOB_NAME);

    network->markOutput(*sigmoid->getOutput(0));

    builder->setMaxBatchSize(maxBatchSize);
    config->setMaxWorkspaceSize(200 * (1 << 20));
    ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);

    network->destroy();

    return engine;


}


void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream) {
    IBuilder *builder = createInferBuilder(gLogger);
    IBuilderConfig *config = builder->createBuilderConfig();

    ICudaEngine *engine = createEngine(maxBatchSize, builder, config, nvinfer1::DataType::kFLOAT);
    assert(engine != nullptr);

    (*modelStream) = engine->serialize();

    engine->destroy();
    builder->destroy();
    config->destroy();
}
void doInference(IExecutionContext &context, float *input, float *output, int batchSize) {
    const ICudaEngine &engine = context.getEngine();

    assert(engine.getNbBindings() == 2);
    void *buffers[2];

    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);

    CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_W * INPUT_H * sizeof(float)));
    CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_W * INPUT_H * sizeof(float),
                          cudaMemcpyHostToDevice, stream));
    context.enqueue(batchSize, buffers, stream, nullptr);
    CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
                          stream));
    cudaStreamSynchronize(stream);

    cudaStreamDestroy(stream);
    CHECK(cudaFree(buffers[inputIndex]));
    CHECK(cudaFree(buffers[outputIndex]));
}

int main() {
    char *trtModelStream{nullptr};
    size_t size{0};

    IHostMemory *modelStream{nullptr};
    APIToModel(1, &modelStream);
    assert(modelStream != nullptr);

    std::ofstream p("crnn.engine");
    if (!p) {
        std::cerr << "can not open plan output file" << std::endl;
        return -1;
    }

    p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
    modelStream->destroy();


    std::ifstream file("crnn.engine", std::ios::binary);
    if (file.good()) {
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream = new char[size];
        assert(trtModelStream);
        file.read(trtModelStream, size);
        file.close();
    } else {
        return -1;
    }

    static float data[3 * INPUT_H * INPUT_W];
    for (int i = 0; i < INPUT_H * INPUT_W * 3; i++) {
        data[i] = 1;

    }

    IRuntime *runtime = createInferRuntime(gLogger);
    assert(runtime != nullptr);
    ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
    assert(engine != nullptr);
    IExecutionContext *context = engine->createExecutionContext();
    assert(context != nullptr);

    float prob[OUTPUT_SIZE];
    doInference(*context, data, prob, 1);
    auto start = std::chrono::system_clock::now();
    for (int i = 0; i < 100; i++) {

        doInference(*context, data, prob, 1);

    }
    auto end = std::chrono::system_clock::now();

    std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << " ms" << std::endl;
    context->destroy();
    engine->destroy();
    runtime->destroy();

//    std::cout << "\n OUTPUT:\n";
//    for (unsigned int i = 0; i < OUTPUT_SIZE; i++) {
//        std::cout << prob[i] << ",";
//        if ((i + 1) % 640 == 0)std::cout << std::endl;
//    }
//
//    std::cout << std::endl;

    return 0;
}
